import { PromptTemplate } from 'langchain/prompts';
import { LLMChain, LLMChainInput } from 'langchain/chains';

export class TaskPrioritizationChain extends LLMChain {
  static fromLLM(fields: Omit<LLMChainInput, 'prompt'>): LLMChain {
    const taskPrioritizationTemplate =
      `You are a task prioritization AI tasked with cleaning the formatting of ` +
      `and reprioritizing the following tasks: {task_names}.` +
      ` Consider the ultimate objective of your team: {objective}.` +
      ` Do not remove any tasks. Return the result as a numbered list, like:` +
      ` #. First task` +
      ` #. Second task` +
      ` Start the task list with number {next_task_id}.`;
    const prompt = new PromptTemplate({
      template: taskPrioritizationTemplate,
      inputVariables: ['task_names', 'next_task_id', 'objective'],
    });
    return new TaskPrioritizationChain({ prompt, ...fields });
  }
}
